import open3d as o3d
from sklearn.cluster import KMeans
import torch
import math
import numpy as np
import colorsys
import os.path
import pickle
from torch.autograd import Variable
import random
from PIL import Image
import matplotlib.pyplot as plt


class GroundEqu:

    def gen_ground_pnt_by_equation(self, pnts, equation, format="leftz"):
        res = []
        for pnt in pnts:
            x = pnt[0]
            y = pnt[1]
            z = 0
            if format == "leftz":
                ac, bc, dc = equation
                z = ac * x + bc * y + dc
            if format == "right0":
                Ax, By, Cz, D = equation
                z = -(Ax * x + By * y + D) / Cz
            res.append([x, y, z])
        return res

    def vis_pnt(self, l):
        # o3d.io.write_point_cloud(f"{self.data}{id}.ply", pcd)
        # pcd_load = o3d.io.read_point_cloud(f"{self.data}{id}.ply")

        colors = [(0.5, 0.5, 1), (0.8, 0.8, 0.8), (0.2, 0.2, 0.2)]
        pnts = []
        for i, v in enumerate(l):
            data = np.asarray(v)
            pcd = o3d.geometry.PointCloud()
            pcd.points = o3d.utility.Vector3dVector(data)
            pcd.paint_uniform_color(colors[i])
            pnts.append(pcd)
        o3d.visualization.draw_geometries(pnts)

    def find_groud_points(self, color, pnt):
        ground = []
        non_ground = []
        num_rows = color.shape[0]
        num_cols = color.shape[1]
        for r in range(0, num_rows):
            for c in range(0, num_cols):
                rgb = color[r, c, :]
                dep = pnt[(r, c)]
                isground = self.is_ground(rgb, dep)
                if isground == "ground":
                    ground.append(dep)
                if isground == "object": non_ground.append(dep)
        return ground, non_ground

    def order_dic(self, cnt, cluster):
        total = sum([v for k, v in cnt.items()])
        listname = []
        start = 0
        for key, value in sorted(cnt.items(), key=lambda i: i[1], reverse=True):
            if (start == 0) or (value / total > 0.2):
                listname.append([key, value / total, cluster[key]])
            start += 1
        return listname

    def define_ground(self, point1, point2, point3):
        x1, y1, z1 = point1
        x2, y2, z2 = point2
        x3, y3, z3 = point3
        a1 = x2 - x1
        b1 = y2 - y1
        c1 = z2 - z1
        a2 = x3 - x1
        b2 = y3 - y1
        c2 = z3 - z1
        a = b1 * c2 - b2 * c1
        b = a2 * c1 - a1 * c2
        c = a1 * b2 - b1 * a2
        d = (- a * x1 - b * y1 - c * z1)
        return [a, b, c, d]

    def distance_by_ground(self, equ_right0, pntlist):
        a, b, c, d = equ_right0
        vec3 = np.array([a, b, c])
        down = np.sqrt(np.sum(np.multiply(vec3, vec3)))

        all = np.array([i for i in pntlist])
        rows = all.shape[0]
        vec1 = np.repeat([[a, b, c, d]], rows, axis=0)

        vec2 = np.ones((rows, 4))
        vec2[:, 0:3] = all
        res1 = np.multiply(vec1, vec2)
        up = np.abs(np.sum(res1, axis=1))
        dis = up / down
        threshhold = np.percentile(dis, 80)
        mask = dis < threshhold
        real_ground = all[mask]
        return real_ground, threshhold

    def min_distance_by_equation(self, ground):
        n = len(ground)
        equation = []
        step = 3
        cnt = 0
        for i in range(0, n, step):
            if i + step < n - 1:
                split = int(step / 3)
                g = ground
                area1 = g[i:split + i]
                area2 = g[split + i:2 * split + i]
                area3 = g[2 * split + i:3 * split + i]
                p1 = np.sum(area1, axis=0) / split
                p2 = np.sum(area2, axis=0) / split
                p3 = np.sum(area3, axis=0) / split
                params = self.define_ground(p1, p2, p3)
                if all(params):
                    equation.append(params)
                else:
                    cnt += 1
        equ = np.array(equation)
        ground4 = np.ones((n, 4))
        equ3 = equ[:, :3]
        ground4[:, 0:3] = np.asarray(ground)
        mul_equ = np.sum(np.multiply(equ3, equ3), axis=1)
        denorminator = np.sqrt(mul_equ)
        numerater = np.sum(np.abs(np.matmul(equ, ground4.T)), axis=1) / n
        dis = np.divide(numerater, denorminator)
        i = np.argmin(dis)
        return equ[i]

    def ground_equation_adjust_by_torch(self, equ, ground):
        a, b, c, d = equ
        g = ground
        x = np.ones((len(ground), 3))
        x[:, 0:2] = g[:, 0:2]
        xy_from_np = torch.from_numpy(x)
        z_from_np = torch.from_numpy(g[:, 2:])

        w_from_np = torch.from_numpy(np.array([[-a / c], [-b / c], [-d / c]]))
        w = Variable(w_from_np, requires_grad=True)
        lr = 1e-2
        forward = lambda x: x.mm(w)
        loss = lambda x, y: torch.pow(y - forward(x), 2).sum()

        rows = xy_from_np.shape[0]
        for i in range(0, rows, 10):
            tx = xy_from_np[i:i + 10, :]
            ty = z_from_np[i:i + 10, :]
            loss_v = loss(tx, ty)
            loss_v.backward()
            with torch.no_grad():
                w -= lr * w.grad
                w.grad.zero_()
        return w.detach().numpy()

    def main(self, groundPoints):

        equ1 = self.min_distance_by_equation(groundPoints)
        # real_ground, threashhold = self.distance_by_ground(equ1, groundPoints)
        # equ2 = self.ground_equation_adjust_by_torch(equ1, real_ground)
        # ground_by_equ=self.gen_ground_pnt_by_equation(groundPoints,equ2,"leftz")
        # ground_by_equ=self.gen_ground_pnt_by_equation(groundPoints,equ1,"right0")
        # self.vis_pnt([ground_by_equ,groundPoints])
        return equ1


if __name__ == '__main__':
    f = GroundEqu()
    d = np.asarray([
        [0.14355469, -1.06134158, 2.60449219],
        [-0.19628906, -0.75292969, 2.3249733],
        [-0.54199219, -0.80834253, 1.35644531],
        [0.14355469, -1.06134158, 2.60449219],
        [-0.19628906, -0.75292969, 2.3249733],
        [-0.54199219, -0.80834253, 1.35644531],
        [0.14355469, -1.06134158, 2.60449219],
        [-0.19628906, -0.75292969, 2.3249733],
        [-0.54199219, -0.80834253, 1.35644531],
    ])

    f.main(d)
